{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jsonlines\n",
    "import random\n",
    "from pprint import pprint\n",
    "from tqdm.auto import tqdm\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(fname):\n",
    "    with jsonlines.open(fname) as reader:\n",
    "        return list(reader)\n",
    "data = load_data('data/codesearchnet_javascript/javascript_train_supervised.jsonl')\n",
    "\n",
    "def extract_fields(fields, gen):\n",
    "    for item in gen:\n",
    "        yield \n",
    "\n",
    "def parse_data(key, data):\n",
    "    matches = filter(lambda x: key.lower() == x['func_name'].lower(), data)\n",
    "    extracted_matches = map(lambda item: {k: v for k, v in item.items() if k in ['func_name', 'code']}, matches)\n",
    "    return list(extracted_matches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "keys = ['sort', 'download', 'search', 'clone', 'bundle', 'destroy', 'query', 'shuffle', 'append', 'log', 'forEach']\n",
    "matches = {}\n",
    "for key in keys:\n",
    "    matches[key] = parse_data(key, data)\n",
    "n_matches = sum(map(len, matches.values()))\n",
    "# matches['negatives'] = random.sample(data, n_matches)\n",
    "\n",
    "for k, v in matches.items():\n",
    "    print(k, len(v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('data/codesearchnet_javascript/tsne_programs.pickle', 'wb') as f:\n",
    "    pickle.dump(matches, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
